import os
from operator import itemgetter

from dotenv import load_dotenv
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDatabaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import ChatOpenAI

load_dotenv()

# 1.创建模型
model = ChatOpenAI(
    model='qwen-plus',
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# SQLAIChemy 数据库工具

HOST_NAME = '127.0.0.1'
PORT = '3306'
# DATABASE = 'ruoyi-vue-pro'
DATABASE = 'wx'
USERNAME = 'root'
PASSWORD = 'root123456'
# mysqlclient 驱动URL
MYSQL_URI = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOST_NAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URI)

# print(db.get_usable_table_names())
# print(db.run('select * from trade_order_item limit 10'))


# 直接使用大模型和数据库整合
# test_chain = create_sql_query_chain(model, db)
# # resp = test_chain.invoke({'question': '请问：订单表中有多少条记录？'})
# # 只能根据你的问题生成SQL语句，并不能执行SQL
# resp = test_chain.invoke({'question': '请问：菜单表中有多少条数据？'})
# print(resp)
# # SQLQuery: `SELECT COUNT(*) AS total_rows FROM sys_menu;`


answer_prompt = PromptTemplate.from_template(
    """给定以下用户问题、SQL语句和SQL执行后的结果，回答用户问题。
    Question: {question}
    SQL Query: {query}
    SQL Result: {result}
    回答："""
)


# Add this processor to clean the SQL output
def clean_sql_output(response):
    """
    清理SQL查询输出的综合函数：
    1. 去除`SQLQuery:`前缀
    2. 去除开头结尾的反引号
    3. 确保最终是合法SQL
    """
    if isinstance(response, str):
        sql = response
    elif isinstance(response, dict) and 'result' in response:
        sql = response['result']
    else:
        return response

    # 清理步骤
    sql = sql.replace("SQLQuery:", "").strip()  # 去除前缀
    sql = sql.strip('`')  # 去除反引号

    # 确保以分号结尾
    if not sql.endswith(';'):
        sql += ';'

    # 返回与输入相同的格式
    if isinstance(response, str):
        return sql
    elif isinstance(response, dict):
        return {**response, 'result': sql}
    return sql


# 创建一个执行SQL语句的工具
execute_sql_tool = QuerySQLDatabaseTool(db=db)

# 生成SQL的chain
sql_query_chain = create_sql_query_chain(model, db) | RunnableLambda(clean_sql_output)

# 生成SQL -》 执行SQL -》
chain = (RunnablePassthrough.assign(query=sql_query_chain).assign(result=itemgetter('query') | execute_sql_tool)
         | answer_prompt
         | model
         | StrOutputParser()
         )

resp = chain.invoke({'question': '请问：菜单表中有多少条数据？'})
print(resp)
# 菜单表中共有 **63 条数据**。

